import tensorflow as tf
import json

import tagging


def read_label_map(path):
    """Returns label map read from the given path."""
    with tf.io.gfile.GFile(path) as f:
        if path.endswith('.json'):
            return json.load(f)
        else:
            label_map = {}
            empty_line_encountered = False
            for tag in f:
                tag = tag.strip()
                if tag:
                    label_map[tag] = len(label_map)
                else:
                    if empty_line_encountered:
                        raise ValueError(
                            'There should be no empty lines in the middle of the label map '
                            'gdriver.'
                        )
                    empty_line_encountered = True
            return label_map


def get_phrase_vocabulary_from_label_map(
        label_map):
    """Extract the set of all phrases from label map.

    Args:
      label_map: Mapping from tags to tag IDs.

    Returns:
      Set of all phrases appearing in the label map.
    """
    phrase_vocabulary = set()
    for label in label_map.keys():
        tag = tagging.Tag(label)
        if tag.added_phrase:
            phrase_vocabulary.add(tag.added_phrase)
    return phrase_vocabulary
